import os
import time
import librosa
import numpy as np
import psutil
from flask import Flask, request, jsonify
from werkzeug.utils import secure_filename

app = Flask(__name__)
UPLOAD_FOLDER = "uploads"
os.makedirs(UPLOAD_FOLDER, exist_ok=True)
app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER


def detect_longest_speech_duration(y, sr, silence_threshold=5, frame_length=1024, hop_length=512):
    frame_amplitudes = np.abs(librosa.stft(y, n_fft=frame_length, hop_length=hop_length))
    frame_max_amplitude = np.max(frame_amplitudes, axis=0)
    amplitude_threshold = max(1, 0.02 * np.max(frame_max_amplitude))

    speech_durations = []
    current_duration = 0
    silence_counter = 0

    for amplitude in frame_max_amplitude:
        if amplitude > amplitude_threshold:
            if silence_counter >= silence_threshold:
                speech_durations.append(current_duration)
                current_duration = 0
            current_duration += 1
            silence_counter = 0
        else:
            silence_counter += 1

    if current_duration > 0:
        speech_durations.append(current_duration)

    return max(speech_durations, default=0) * hop_length / sr


def analyze_audio(file_path):
    process = psutil.Process()
    start_memory = process.memory_info().rss / 1024 / 1024
    start_time = time.time()

    y, sr = librosa.load(file_path, sr=None)
    longest_duration = detect_longest_speech_duration(y, sr)

    f0, _, _ = librosa.pyin(y, fmin=80, fmax=250)
    f0_cleaned = f0[~np.isnan(f0)]
    jitter = np.mean(np.abs(np.diff(f0_cleaned)) / f0_cleaned[:-1]) * 100 if len(f0_cleaned) > 1 else 0

    amplitude = np.abs(y)
    amplitude_threshold = 0.001
    amplitude_filtered = amplitude[amplitude > amplitude_threshold]
    amplitude_jitter = np.mean(np.abs(np.diff(amplitude_filtered)) / amplitude_filtered[:-1]) * 100 if len(
        amplitude_filtered) > 1 else 0

    mean_f0 = np.mean(f0_cleaned) if len(f0_cleaned) > 0 else 0
    max_f0 = np.max(f0_cleaned) if len(f0_cleaned) > 0 else 0
    min_f0 = np.min(f0_cleaned) if len(f0_cleaned) > 0 else 0
    mean_amplitude = np.mean(amplitude_filtered) if len(amplitude_filtered) > 0 else 0

    dsi = 0.13 * longest_duration + 0.0053 * max_f0 - 0.26 * 35 - 1.18 * jitter * 0.01 + 12.4
    end_time = time.time()

    # Convert all NumPy float32 to Python float
    return {
        "最长持续发声时间": round(float(longest_duration), 2),
        "基频抖动 (Jitter)": round(float(jitter), 4),
        "振幅抖动 (Amplitude Jitter)": round(float(amplitude_jitter), 4),
        "平均基频 (Mean F0)": round(float(mean_f0), 2),
        "最小基频 (Min F0)": round(float(min_f0), 2),
        "最大基频 (Max F0)": round(float(max_f0), 2),
        "平均振幅 (Mean Amplitude)": round(float(mean_amplitude), 4),
        "DSI": round(float(dsi), 4),
    }



@app.route("/analyze", methods=["POST"])
def upload_file():
    if 'file' not in request.files:
        return jsonify({"error": "No file uploaded"}), 400

    file = request.files['file']
    if file.filename == '':
        return jsonify({"error": "No selected file"}), 400

    filename = secure_filename(file.filename)
    file_path = os.path.join(app.config['UPLOAD_FOLDER'], filename)
    file.save(file_path)

    result = analyze_audio(file_path)
    os.remove(file_path)

    return jsonify(result)


if __name__ == "__main__":
    app.run(host="0.0.0.0", port=8001, debug=True)
